#ifndef COMMON_LOGGER_HPP
#define COMMON_LOGGER_HPP

#include <iostream>
#include <fstream>
#include <sstream>
#include <mutex>
#include <string>
#include <ctime>
#include <memory>
#include <mpi.h>
#include <vector>
#include <algorithm>
#include <chrono>
#include <iomanip>

namespace common {

enum class LogLevel {
    TRACE,
    DEBUG,
    INFO,
    WARN,
    ERROR
};

class Logger {
public:
    static Logger& getInstance() {
        static Logger instance;
        return instance;
    }

    void init(LogLevel level = LogLevel::INFO, bool console_output = true, 
              const std::string& file_prefix = "log") {
        std::lock_guard<std::mutex> lock(mutex_);
        level_ = level;
        console_output_ = console_output;
        
        // Get MPI rank
        int rank;
        MPI_Comm_rank(MPI_COMM_WORLD, &rank);
        rank_ = rank;

        // Open log file if prefix is provided
        if (!file_prefix.empty()) {
            std::string filename = file_prefix + "_rank" + std::to_string(rank_) + ".log";
            file_.open(filename, std::ios::out | std::ios::app);
        }
    }

    template<typename... Args>
    void log(LogLevel level, const char* file, int line, Args... args) {
        if (level < level_) return;

        std::lock_guard<std::mutex> lock(mutex_);
        std::stringstream ss;
        
        // Add timestamp with millisecond precision
        auto now = std::chrono::system_clock::now();
        auto now_time_t = std::chrono::system_clock::to_time_t(now);
        auto now_ms = std::chrono::duration_cast<std::chrono::milliseconds>(
            now.time_since_epoch()
        ).count() % 1000;
        
        char timestamp[32];
        std::strftime(timestamp, sizeof(timestamp), "%Y-%m-%d %H:%M:%S", std::localtime(&now_time_t));
        
        ss << "[" << timestamp << "." << std::setfill('0') << std::setw(3) << now_ms << "] "
           << "[Rank " << rank_ << "] "
           << "[" << getLevelString(level) << "] "
           << "[" << file << ":" << line << "] ";
        
        (ss << ... << args) << std::endl;

        // Output to console if enabled
        if (console_output_) {
            std::cout << ss.str();
        }

        // Output to file if opened
        if (file_.is_open()) {
            file_ << ss.str();
            file_.flush();
        }
    }

    static void mergeLogFiles(const std::string& file_prefix, const std::string& output_file) {
        int rank, size;
        MPI_Comm_rank(MPI_COMM_WORLD, &rank);
        MPI_Comm_size(MPI_COMM_WORLD, &size);

        // Only rank 0 performs the merge
        if (rank != 0) return;

        struct LogEntry {
            std::string timestamp;
            std::string content;
            bool operator<(const LogEntry& other) const {
                return timestamp < other.timestamp;
            }
        };

        std::vector<LogEntry> all_entries;

        // Read all log files
        for (int i = 0; i < size; ++i) {
            std::string filename = file_prefix + "_rank" + std::to_string(i) + ".log";
            std::ifstream input(filename);
            
            if (!input.is_open()) {
                std::cerr << "Warning: Could not open " << filename << std::endl;
                continue;
            }

            std::string line;
            while (std::getline(input, line)) {
                // Extract timestamp from line
                // Format: [2024-03-14 10:30:15] [other content...]
                if (line.length() < 21) continue; // Skip invalid lines
                
                LogEntry entry;
                entry.timestamp = line.substr(1, 19); // Extract timestamp
                entry.content = line;
                all_entries.push_back(entry);
            }
            input.close();
        }

        // Sort entries by timestamp
        std::sort(all_entries.begin(), all_entries.end());

        // Write merged output
        std::ofstream output(output_file);
        if (!output.is_open()) {
            std::cerr << "Error: Could not open output file " << output_file << std::endl;
            return;
        }

        for (const auto& entry : all_entries) {
            output << entry.content << std::endl;
        }
        output.close();
    }

    ~Logger() {
        if (file_.is_open()) {
            file_.close();
        }
    }

private:
    Logger() = default;
    Logger(const Logger&) = delete;
    Logger& operator=(const Logger&) = delete;

    const char* getLevelString(LogLevel level) {
        switch (level) {
            case LogLevel::TRACE: return "TRACE";
            case LogLevel::DEBUG: return "DEBUG";
            case LogLevel::INFO:  return "INFO";
            case LogLevel::WARN:  return "WARN";
            case LogLevel::ERROR: return "ERROR";
            default: return "UNKNOWN";
        }
    }

    std::mutex mutex_;
    LogLevel level_{LogLevel::INFO};
    bool console_output_{true};
    std::ofstream file_;
    int rank_{0};
};

// Convenience macros for logging
#define LOG_TRACE(...) \
    common::Logger::getInstance().log(common::LogLevel::TRACE, __FILE__, __LINE__, __VA_ARGS__)
#define LOG_DEBUG(...) \
    common::Logger::getInstance().log(common::LogLevel::DEBUG, __FILE__, __LINE__, __VA_ARGS__)
#define LOG_INFO(...) \
    common::Logger::getInstance().log(common::LogLevel::INFO, __FILE__, __LINE__, __VA_ARGS__)
#define LOG_WARN(...) \
    common::Logger::getInstance().log(common::LogLevel::WARN, __FILE__, __LINE__, __VA_ARGS__)
#define LOG_ERROR(...) \
    common::Logger::getInstance().log(common::LogLevel::ERROR, __FILE__, __LINE__, __VA_ARGS__)

} // namespace common

#endif // COMMON_LOGGER_HPP 